from sklearn.cluster import MiniBatchKMeans


def run(df, params):
    df = df.dropna()

    k = int(params["k"])
    max_iter = int(params["max_iter"])

    mlmodel = MiniBatchKMeans(n_clusters=k, max_iter=max_iter, batch_size=10000).fit(df[params["cols"]].values)
    pred = mlmodel.labels_

    output_cols = "_cluster_id_"
    df[output_cols] = pred.astype("int")

    return df
